import os
import xml.etree.ElementTree as ET
import cv2
from PIL import Image
from tqdm import tqdm
import numpy as np

import tensorflow as tf

from utils_map import get_map
from research.object_detection.utils import label_map_util
from research.object_detection.utils import visualization_utils as viz_utils

gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

if __name__ == '__main__':

    #   map_mode为0代表整个map计算流程，包括获得预测结果、获得真实框、计算VOC_map。
    #   map_mode为1代表仅仅获得预测结果。
    #   map_mode为2代表仅仅获得真实框。
    #   map_mode为3代表仅仅计算VOC_map。

    map_out_path = './map_out'

    map_mode = 0

    MINOVERLAP = 0.1

    label_path = 'E:/Python/models/research/object_detection/mydata/trash_label_map.pbtxt'
    path = 'E:/Python/models/research/object_detection/mydata'
    path_saved_model = path + "/saved_model"
    detect_fn = tf.saved_model.load(path_saved_model)
    category_index = label_map_util.create_category_index_from_labelmap(label_path, use_display_name=True)

    image_ids = open(os.path.join(path, "test.txt")).read().strip().split()


    if not os.path.exists(map_out_path):
        os.makedirs(map_out_path)
    if not os.path.exists(os.path.join(map_out_path, 'ground-truth')):
        os.makedirs(os.path.join(map_out_path, 'ground-truth'))
    if not os.path.exists(os.path.join(map_out_path, 'detection-result')):
        os.makedirs(os.path.join(map_out_path, 'detection-result'))

    if map_mode == 0 or map_mode == 1:
        print('Get Predict.')
        for image_id in tqdm(image_ids):
            image_path = os.path.join(path, "JPEGImages/" + image_id + ".jpg")
            image_np = cv2.imread(image_path)

            input_tensor = tf.convert_to_tensor(image_np)
            input_tensor = input_tensor[tf.newaxis, ...]
            detections = detect_fn(input_tensor)
            num_detections = int(detections.pop('num_detections'))
            detections = {key: value[0, :num_detections].numpy()
                          for key, value in detections.items()}
            detections['num_detections'] = num_detections

            # detection_classes should be ints.
            detections['detection_classes'] = detections['detection_classes'].astype(np.int64)

            img_height, img_width = image_np.shape[:2]

            image_np_with_detections = image_np.copy()

            with open(os.path.join(map_out_path, "detection-result/" + image_id + ".txt"), "w") as new_f:

                viz_utils.visualize_boxes_and_labels_on_image_array(
                    image_np_with_detections,
                    detections['detection_boxes'],
                    detections['detection_classes'],
                    detections['detection_scores'],
                    category_index,
                    use_normalized_coordinates=True,
                    max_boxes_to_draw=500,
                    min_score_thresh=0.1,
                    skip_scores=False,
                    agnostic_mode=False)

                class_name = viz_utils.trash_class
                num_detect = viz_utils.num_detection
                location = viz_utils.location_list
                score = viz_utils.score

                for i in range(len(class_name)):
                    # print('trash:', class_name[i])
                    # print('score', round(score[i], 4))
                    # print('location:', location[i][0], location[i][1], location[i][2], location[i][3])
                    new_f.write("%s %s %s %s %s %s\n" % (
                        class_name[i], round(score[i], 4), int(location[i][1] * img_width), int(location[i][0] * img_height), int(location[i][3] * img_width), int(location[i][2] * img_height)))
        print('Get Predict Done.')

    if map_mode == 0 or map_mode == 2:
        print('Get Ground Truth.')
        for image_id in tqdm(image_ids):
            with open(os.path.join(map_out_path, "ground-truth/" + image_id + ".txt"), 'w') as new_f:
                root = ET.parse(os.path.join(path, "Annotations/" + image_id + ".xml")).getroot()
                for obj in root.findall('object'):
                    difficult_flag = False
                    if obj.find('difficult') != None:
                        difficult = obj.find('difficult').text
                        if int(difficult) == 1:
                            difficult_flag = True
                    obj_name = obj.find('name').text
                    if obj_name not in ['trash1', 'trash2']:
                        continue
                    bndbox = obj.find('bndbox')
                    left = bndbox.find('xmin').text
                    top = bndbox.find('ymin').text
                    right = bndbox.find('xmax').text
                    bottom = bndbox.find('ymax').text

                    if difficult_flag:
                        new_f.write("%s %s %s %s %s difficult\n" % (obj_name, left, top, right, bottom))
                    else:
                        new_f.write("%s %s %s %s %s\n" % (obj_name, left, top, right, bottom))
        print("Get ground truth result done.")

    if map_mode == 0 or map_mode == 3:
        print("Get map.")
        get_map(MINOVERLAP, True, path=map_out_path)
        print("Get map done.")
